"""
@description: Dataset loader class
@author: Zzay
@create: 2022/07/08
"""
import numpy as np
import cv2
import os


class SimpleDatasetLoader:
    """
    Dataset loader class.
    """

    def __init__(self, preprocessors=None):
        """
        Store the image preprocessors.
        :param preprocessors: Preprocessors
        """
        # store the image preprocessor
        self.preprocessors = preprocessors
        # if the preprocessors are None, initialize them as an empty list
        if self.preprocessors is None:
            self.preprocessors = []

    def load(self, image_paths, verbose=500, grayscale=False):
        """
        Load datasets.
        """
        # initialize the list of features and labels
        data = []
        labels = []
        # loop over the input images
        for (i, image_path) in enumerate(image_paths):
            # load the image and extract the class label (path format: "/path/to/dataset/{class}/{image}.jpg)
            image = cv2.imread(image_path)
            if grayscale:
                image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            label = image_path.split(os.path.sep)[-2]
            # check if the preprocessors are None or not
            if self.preprocessors is not None:
                # loop over the preprocessors and apply each to the image
                for preprocessor in self.preprocessors:
                    image = preprocessor.preprocess(image)
            # treat the processed image as a "feature vector" by updating the data list followed by the labels
            data.append(image)
            labels.append(label)
            # show an update every `verbose` images
            if verbose > 0 and i > 0 and (i + 1) % verbose == 0:
                print("[INFO] processed {}/{}".format(i + 1, len(image_paths)))

        # return a tuple of the data and labels
        return np.array(data), np.array(labels)
